from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

from zhipuai import ZhipuAI
client = ZhipuAI()
model = "cogvideox"

def get_completion(model="cogvideox", 
                   prompt=None, 
                   image_url=None):
    response = client.videos.generations(
      model=model, 
      prompt=prompt,
      image_url=image_url
    )
    return response

while True:  
    user_input = input("请输入待生成视频的描述（输入'exit'退出程序）: ")  
    if user_input.lower() == 'exit':  
        print("程序退出。\n")  
        exit()

    image_url = input("请输入待生成视频的基础图片（输入'exit'退出程序）: ") 
    if image_url.lower() == 'exit':  
        print("程序退出。\n")  
        exit()

    if user_input == None and image_url == None:
        print("描述和基础图片必须输入至少一项。") 
        continue
    else:
        break

res = get_completion(model=model, 
                     prompt=user_input, 
                     image_url=image_url)  
print("===提交任务===") 
print(res)

print("===查询任务状态===")     
while True:
    response = client.videos.retrieve_videos_result(
        id=f"{res.id}"
    )
    if response.task_status == "PROCESSING":
        print("PROCESSING...")
        import time
        time.sleep(30)
        continue
    elif response.task_status == "FAIL":
        print("===任务失败===")
        break
    else:
        print("===任务完成===")
        print(f"视频URL：{response.video_result[0].url}")
        print(f"视频封面：{response.video_result[0].cover_image_url}\n")
        break